#!/usr/bin/python
# -*- coding: utf-8 -*-
from pandas import DataFrame
from numpy import nan

from core.index_base import IndexBase


class DC(IndexBase):
    """
    Donchian Channel 唐奇安通道
    """
    def __init__(self, period = 20):
        super(DC, self).__init__()
        self.period = period

    def compute(self, data:DataFrame):
        dc_upbound_key = 'dc_upbound' + str(self.period)
        dc_downbound_key = 'dc_downbound' + str(self.period)
        dc_midbound_key = 'dc_midbound' + str(self.period)
        self.added_keys.extend([dc_upbound_key, dc_midbound_key, dc_downbound_key])
        data[dc_upbound_key] = nan
        data[dc_downbound_key] = nan
        data[dc_midbound_key] = nan
        data[dc_upbound_key] = data['high'].rolling(self.period).max()
        data[dc_downbound_key] = data['low'].rolling(self.period).min()
        data[dc_midbound_key] = (data[dc_upbound_key] + data[dc_downbound_key]) / 2
